package Trie.Medium;

import java.util.*;

public class LC1268 {
    private class TrieNode {
        TrieNode[] children = new TrieNode[26];
        TreeSet<String> prefixWords = new TreeSet<>(); // 以该结点为前缀的产品集合
    }

    private class Trie {
        TrieNode root = new TrieNode();

        public Trie() {
        }

        public void insert(String word) {
            char[] charArray = word.toCharArray();
            TrieNode ptr = root;
            for (int i = 0; i < charArray.length; i++) {
                char ch = charArray[i];
                if (ptr.children[ch - 'a'] == null) ptr.children[ch - 'a'] = new TrieNode();
                ptr = ptr.children[ch - 'a'];
                ptr.prefixWords.add(word);
            }
        }

        public TreeSet<String> query(String prefix) {
            TrieNode ptr = root;
            char[] charArray = prefix.toCharArray();
            for (int i = 0; i < charArray.length; i++) {
                char ch = charArray[i];
                if (ptr.children[ch - 'a'] == null) return null;
                ptr = ptr.children[ch - 'a'];
            }
            return ptr.prefixWords;
        }
    }

    public List<List<String>> suggestedProducts(String[] products, String searchWord) {
        Trie trie = new Trie();
        for (String product : products) trie.insert(product);

        TrieNode ptr = trie.root;
        List<List<String>> ret = new ArrayList<>();
        char[] charArray = searchWord.toCharArray();
        for (int i = 0; i < charArray.length; i++) {
            char ch = charArray[i];
            if (ptr.children[ch - 'a'] == null) {
                for (int j = 0; j < charArray.length - i; j++) ret.add(new ArrayList<>());
                break;
            }
            ptr = ptr.children[ch - 'a'];
            List<String> list = new ArrayList<>(ptr.prefixWords);
            Collections.sort(list);
            ret.add(list.subList(0, Math.min(3, list.size())));
        }
        return ret;
    }
}
